import os
import torch
from torch.utils.data import Dataset
from PIL import Image
import h5py


# utils
origin_datadir = r'/home/yeadc/Documents/cxj/datasets_classify_v2'  # 初始数据集位置
categories = ['ADI_particle_developed', 'Array_peeling', 'Cu_missing', 'Other_peeling', 'Partial_etch', 'Pattern_fail',
              'PR_peeling', 'Seam', 'Reference', 'Surface_particle', 'Burried_particle', 'Cu_diffuse', 'Prelayer_defect_developed',
              'Void', 'Residue', 'Scratch']

categories_to_id = dict((c, i) for i, c in enumerate(categories))
id_to_categories = dict((v, k) for k, v in categories_to_id.items())


def get_imgs_labels(dir):
    # dir: E:\proj\AI\dataset\build_dataset\datasets\train
    imgs = []
    labels = []
    categories = os.listdir(dir)
    for category in categories:
        category_path = os.path.join(dir, category)
        imgs_path = os.listdir(category_path)
        for img in imgs_path:
            imgs.append(os.path.join(category_path, img))
            labels.append(categories_to_id[category])
    return imgs, labels

class Mydataset(Dataset):
    # 类初始化
    def __init__(self, img_paths, labels, transform_class, transform_defect):
        self.img_paths = img_paths
        self.labels = labels
        self.transform_class = transform_class
        self.transform_defect = transform_defect

    # 进行切片
    def __getitem__(self, index): #根据给出的索引进行切片，并对其进行数据处理转换成Tensor，返回成Tensor
        img_path = self.img_paths[index]
        label = self.labels[index]
        # get images from h5file
        f = h5py.File(img_path, 'r')  # 打开h5文件
        class_tiff = f['class_tiff'][:]
        defect_tiff = f['defect_tiff'][:]
        reference_tiff = f['reference_tiff'][:]
        f.close()


        class_tiff = Image.fromarray(class_tiff)
        defect_tiff = Image.fromarray(defect_tiff)
        reference_tiff = Image.fromarray(reference_tiff)

        if self.transform_class is not None:
            class_tiff = self.transform_class(class_tiff)

        if self.transform_defect is not None:
            defect_tiff = self.transform_defect(defect_tiff)
            reference_tiff = self.transform_defect(reference_tiff)
            # print(f'image size after transform:{image.shape}')#torch.Size([1, 256, 256])
        sample = {'class_tiff': class_tiff, 'defect_tiff': defect_tiff, 'reference_tiff': reference_tiff, 'label': label}
        return sample

    # 返回长度
    def __len__(self):
        return len(self.img_paths)

# class Mydataset(torch.utils.data.Dataset):
#     # 类初始化
#     def __init__(self, img_paths, labels, transform):
#         self.img_paths = img_paths
#         self.labels = labels
#         self.transform = transform
#     # 进行切片
#     def __getitem__(self, index): #根据给出的索引进行切片，并对其进行数据处理转换成Tensor，返回成Tensor
#         img_path = self.img_paths[index]
#         label = self.labels[index]
#         image = Image.open(img_path)
#         # image = cv2.imread(img_path, 1)
#         # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
#         # image = Image.fromarray(image)
#         if self.transform is not None:
#             image = self.transform(image)
#             # print(f'image size after transform:{image.shape}')#torch.Size([1, 256, 256])
#         return image, label
#     # 返回长度
#     def __len__(self):
#         return len(self.img_paths)